"""
Author: Morphlng
Date: 2024-04-30 13:36:05
LastEditTime: 2024-12-14 21:30:47
LastEditors: Morphlng
Description: Stub action space for testing
FilePath: /DrivingGym/src/driving_gym/environment/agent/action/stub_action.py
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np

from driving_gym.environment.agent.action.action_interface import (
    AbstractAction,
    ActionInterface,
)
from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface


class PseudoAction(AbstractAction):
    def __init__(self, action, duration=1):
        super().__init__(action, duration)

    @override(AbstractAction)
    def run_step(self):
        """Return a carla control signal dict."""
        return self.to_dict()


class StubAction(ActionInterface):
    def __init__(self, action_config: dict, adapter: AdapterInterface):
        """Initialize the action converter for walker action space

        Args:
            action_config (dict): A dictionary of action config
        """
        super().__init__(action_config, adapter)

        self._mask_on = action_config.get("mask_on", False)
        self._is_discrete = action_config.get("is_discrete", False)
        self._multi_discrete = action_config.get("multi_discrete", False)
        if self._is_discrete:
            self._action_set = action_config.get("action_set", {i: i for i in range(4)})
        elif self._multi_discrete:
            self._action_shape = action_config.get(
                "action_set", ({i: i for i in range(4)}, {i: i for i in range(3)})
            )
        else:
            self._action_shape = action_config.get("action_shape", (2,))

    @override(ActionInterface)
    def convert_single_action(self, action):
        """Convert the action of a model output to an AbstractAction instance

        Args:
            action: Action input from the model

        Returns:
            DirectAction: A direct action instance
        """
        if self._is_discrete:
            action = self._action_set[action]
        elif self._multi_discrete:
            action = (self._action_shape[i][action[i]] for i in range(len(action)))
        return PseudoAction(action)

    @override(ActionInterface)
    def get_action_space(self):
        if self._is_discrete:
            return gym.spaces.Discrete(len(self._action_set))
        elif self._multi_discrete:
            return gym.spaces.MultiDiscrete(
                [len(self._action_shape[i]) for i in range(len(self._action_shape))]
            )
        else:
            return gym.spaces.Box(
                low=-1.0, high=1.0, shape=self._action_shape, dtype=np.float32
            )

    @override(ActionInterface)
    def get_mask_space(self):
        if not self._mask_on:
            return None

        return super().get_mask_space()

    @override(ActionInterface)
    def get_action_mask(self) -> np.ndarray:
        if not self._mask_on:
            return None

        if self._is_discrete:
            return np.random.choice([0, 1], size=len(self._action_set), p=[0.5, 0.5]).astype(np.int8)  # fmt: skip
        elif self._multi_discrete:
            mask = []
            for i in range(len(self._action_shape)):
                mask += np.random.choice([0, 1], size=len(self._action_shape[i]), p=[0.5, 0.5]).tolist()  # fmt: skip
            return np.array(mask, dtype=np.int8)
        else:
            raise ValueError(f"Invalid action space for mask")
